# Flatten
# 将卷积层传过来的特征图按像素点拉直，拼接成一个向量，其分量个数=通道数x高x宽
import numpy as np
from Layer import Layer


class Flatten(Layer):
    def __init__(self):
        super().__init__()
        self.input_shape = []

    # 矩阵变向量
    def forward(self, input_tensor):
        self.input_shape = input_tensor.shape
        input_tensor = input_tensor.reshape(input_tensor.shape[0], -1)
        return input_tensor

    # 向量变矩阵
    def backward(self, error_tensor):
        error_tensor = error_tensor.reshape(*self.input_shape)
        return error_tensor
